use clap::{Parser, Subcommand};
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    fs,
    path::{Path, PathBuf},
    process::{Command, Stdio},
};

#[derive(Parser)]
#[command(name = "yam")]
#[command(version = "1.0")]
#[command(about = "YAM - Yet Another Model Manager")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Pull a model from repository
    Pull { model: String },
    /// Run a model
    Run { model: String },
    /// Remove a model
    Rm { model: String },
    /// List downloaded models
    List,
    /// List running models
    Ps,
    /// Show model details
    Show { model: String },
}

#[derive(Serialize, Deserialize, Debug)]
struct ModelMetadata {
    name: String,
    version: String,
    size: u64,
    description: String,
    format: String,
    parameters: HashMap<String, String>,
}

struct ModelManager {
    models_dir: PathBuf,
}

impl ModelManager {
    fn new() -> Self {
        let home_dir = dirs::home_dir().expect("Failed to get home directory");
        let models_dir = home_dir.join(".yam").join("models");
        fs::create_dir_all(&models_dir).expect("Failed to create models directory");
        Self { models_dir }
    }

    fn get_model_path(&self, model: &str) -> PathBuf {
        self.models_dir.join(model)
    }

    fn model_exists(&self, model: &str) -> bool {
        self.get_model_path(model).exists()
    }

    fn list_models(&self) -> Vec<String> {
        fs::read_dir(&self.models_dir)
            .unwrap()
            .filter_map(|entry| {
                let entry = entry.ok()?;
                if entry.file_type().ok()?.is_dir() {
                    entry.file_name().into_string().ok()
                } else {
                    None
                }
            })
            .collect()
    }

    fn remove_model(&self, model: &str) -> Result<(), String> {
        let path = self.get_model_path(model);
        if path.exists() {
            fs::remove_dir_all(path).map_err(|e| e.to_string())
        } else {
            Err("Model not found".into())
        }
    }

    fn get_model_metadata(&self, model: &str) -> Result<ModelMetadata, String> {
        let metadata_path = self.get_model_path(model).join("metadata.yaml");
        let content = fs::read_to_string(metadata_path).map_err(|e| e.to_string())?;
        serde_yaml::from_str(&content).map_err(|e| e.to_string())
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();
    let model_manager = ModelManager::new();

    match &cli.command {
        Commands::Pull { model } => {
            println!("Pulling model: {}", model);
            // TODO: Implement actual download logic
            // This would include downloading from a model repository
            // and saving to models directory
            Ok(())
        }

        Commands::Run { model } => {
            if !model_manager.model_exists(model) {
                return Err(format!("Model {} not found", model).into());
            }

            println!("Running model: {}", model);
            // TODO: Implement actual model execution
            // Example using a child process:
            let _child = Command::new("python") // Replace with actual model runner
                .arg("-m")
                .arg("model_runner")
                .arg(model)
                .stdout(Stdio::inherit())
                .spawn()?;
            
            Ok(())
        }

        Commands::Rm { model } => {
            model_manager.remove_model(model)?;
            println!("Removed model: {}", model);
            Ok(())
        }

        Commands::List => {
            let models = model_manager.list_models();
            println!("Downloaded models:");
            for model in models {
                println!("- {}", model);
            }
            Ok(())
        }

        Commands::Ps => {
            println!("Running models:");
            // TODO: Implement process listing logic
            // This would check for running model processes
            Ok(())
        }

        Commands::Show { model } => {
            let metadata = model_manager.get_model_metadata(model)?;
            println!("Model: {}", metadata.name);
            println!("Version: {}", metadata.version);
            println!("Size: {} MB", metadata.size / 1024 / 1024);
            println!("Description: {}", metadata.description);
            println!("Parameters:");
            for (key, value) in &metadata.parameters {
                println!("  {}: {}", key, value);
            }
            Ok(())
        }
    }
}